/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file

    Contains the necessary definitions to use grid factories.
    \ingroup mcrx */

// $Id$

#ifndef __grid_factory__
#define __grid_factory__

#include "grid.h"
#include <vector>
#include "boost/thread/thread.hpp"
#include "boost/thread/mutex.hpp"
#include "counter.h"

namespace mcrx {
  template<typename> class grid_creator;
};


/** Abstract base class for objects which can fill a grid with data.
    Classes deriving from this have to define the methods, which are
    then used by the grid_creator object to build an adaptive grid
    based on the information obtained from the grid factory. 
    \ingroup mcrx */
template <typename cell_data_type>
class mcrx::grid_factory {
public:
  typedef cell_data_type T_data;
  typedef grid_base<T_data> T_grid;
  typedef grid_cell<T_data> T_cell;
  
  /** Predicate indicating whether the specified cell should be
      refined or not. The existing refinement level of the cell is
      also supplied. */
  virtual bool refine_cell_p (const T_cell& c, int level) = 0;
  /** Returns the data for the cell. */
  virtual T_data get_data (const T_cell& c) = 0;
  /** Predicate indicating whether the specified range of cells should
      be unified or not. */
  virtual bool unify_cell_p (typename T_grid::const_local_iterator begin,
			     typename T_grid::const_local_iterator end) = 0;
  /** The number of threads to use for grid building. NOTE that due to
      a bug, this must currently not be greater than one. */
  virtual int n_threads () = 0;
  virtual int work_chunk_levels () {return 0;};
  virtual int estimated_levels_needed () {return 0;};
}; 


/** Grid constructor which uses a grid factory to build the grid.  */
template <typename cell_data_type/*, typename grid_type*/>
mcrx::adaptive_grid<cell_data_type/*, grid_type*/>::
adaptive_grid (const vec3d & mi, const vec3d & ma, const vec3i & nn,
	       grid_factory<T_data>& f):
  g_(mi, ma, nn), subgrid_block (0)
{
  // The grid creator does the work (in the constructor)
  grid_creator<adaptive_grid<cell_data_type/*,grid_type*/> > gc (*this, f);
};


/** Creates a grid using a grid factory to obtain the grid data.
    Grid_creator is a utility object that has the same function as the
    nbody_data_grid when it comes to doing refinement.  (In fact,
    makegrid could have been written to use this mechanism instead,
    but it predates the grid factory functionality.)  \ingroup mcrx */
template <typename grid_type>
class mcrx::grid_creator {
  typedef typename grid_type::T_grid T_grid;
  typedef typename T_grid::T_data T_data;
  typedef typename T_grid::T_cell T_cell;
  typedef grid_factory<T_data> T_factory;
  
  T_factory& factory;
  counter ctr;
  
  /// \name Threading data structures.
  ///@{
  int work_chunk_levels;
  mutable boost::mutex cell_stack_mutex;
  std::vector<std::pair<T_cell*, int> > cell_stack;
  /// Function object used to start threads.
  class thread_start; 
  std::vector<int> creation_stats;
  void pop_cell_and_refine ();
  ///@}

  /// \name Functions for grid building and refinement.
  ///@{
  class refinement_accuracy_data;
  void build_grid (grid_type& g); 
  void create_balanced_queue (T_cell*c, int level);
  refinement_accuracy_data recursive_refine (T_grid& g, int level,
					     std::vector<int>&);
  refinement_accuracy_data recursive_refine_body (T_cell* c, int level,
						  std::vector<int>&);
  void unrefine_if_possible (T_cell& c, refinement_accuracy_data& racc);
  ///@}

public:
  grid_creator (grid_type & g, T_factory & f);
};


/** Keeps information about the grid refinement accuracy when creating
    the grid.  Note how this class is functionally identical to
    mcrx::refinement_accuracy_data. \ingroup mcrx */
template <typename grid_type>
class mcrx::grid_creator<grid_type>::refinement_accuracy_data {
public:
  T_data sum;
  T_data sumsq;
  int n;
  bool all_leaves;

  refinement_accuracy_data (const T_data& s, const T_data& ssq, 
			    int nn, bool al):
    sum (s), sumsq (ssq),
    n (nn), all_leaves (al) {}; 

  // Increment operator is used to accumulate over cells about to be
  // unified, notice that all_leaves is AND'ed together.
  refinement_accuracy_data& operator+= (const refinement_accuracy_data& rhs) {
    sum += rhs.sum; sumsq += rhs.sumsq; n+= rhs.n; all_leaves &= rhs.all_leaves;
    return *this;};
};

template <typename grid_type>
mcrx::grid_creator<grid_type>::
grid_creator (grid_type& g, T_factory& f):
  factory (f), ctr (1) {
  build_grid (g);
}


/** Functor used by boost_threads to start the threads. Simply calls
    pop_cell_and_refine () for the grid factory. \ingroup mcrx */
template <typename grid_type>
class mcrx::grid_creator<grid_type>::thread_start {
private:
  grid_creator<grid_type>* self;
public:
  thread_start (grid_creator<grid_type>* g): self (g) {};
  void operator () () {
    self->pop_cell_and_refine ();
  };
};

  

template <typename grid_type>
typename mcrx::grid_creator<grid_type>::refinement_accuracy_data
mcrx::grid_creator<grid_type>::recursive_refine (T_grid& g, 
					 int level, std::vector<int>& stats)
{
  // We must initialize the structure in this way because of blitz
  // array semantics
  typename T_grid::local_iterator c = g.lbegin();
  refinement_accuracy_data racc (recursive_refine_body (c, level, stats));
  ++c;

  // and now loop over the other cells
  for (; c != g.lend(); ++c) {
    racc+= recursive_refine_body (c, level, stats);
  }

  return racc;
}

template <typename grid_type>
void
mcrx::grid_creator<grid_type>::
unrefine_if_possible (T_cell& c,
		      refinement_accuracy_data& racc)
{
  assert (!c.is_leaf());

  /** \todo THERE IS A PROBLEM HERE

  // How do we combine the new unification scheme with giving the
  // factory latitude to decide whether or not the cells should be
  // unified????  Ask the factory if the cells can be unrefined */

  // (Note that we cannot use begin/end but must use lbegin/lend,
  // because end relies on the NEXT grid cell in the supergrid to not
  // change, and this is not true when we're doing grid creation with
  // multiple threads.)
  T_grid*const g = c.sub_grid();
  typename T_grid::const_local_iterator e =g->lend();
  const bool p =factory.unify_cell_p (g->lbegin(), e);
  assert(e==g->lend()); // is it changing from under us?
  if (p) {
    // yes, unrefine
    // this "constructor" returns a T_data object that represents the unification 
    T_data*const d =  new T_data (T_data::unification (racc.sum, racc.n)) ;
    c.unrefine(d);
  }
  else
    // Mark that the sub cells were not unified
    racc.all_leaves = false;
}

/** This function is called by the thread_start object for each of the
    threads. */
template <typename grid_type>
void 
mcrx::grid_creator<grid_type>::pop_cell_and_refine()
{
  std::cout << "thread starting\n";
  std::vector<int> stats;
  while (true) {

    // pop another cell off the stack, exit if empty
    T_cell* current_cell;
    int start_level;
    {
      // open scope for locking stack_mutex
      boost::mutex::scoped_lock stack_lock (cell_stack_mutex);
      if (!stats.empty()) {
	if (stats.size() > creation_stats.size())
	  creation_stats.resize( stats.size()  );
	for (int i = 0; i < stats.size(); ++i) 
	  creation_stats [i]+= stats [i];
      }
      if (cell_stack.empty())
	break;
      current_cell = cell_stack.back().first;
      start_level  = cell_stack.back().second;
      cell_stack.pop_back();
      ctr++;
      
    }

    recursive_refine_body (current_cell, start_level, stats);
  } // while 

  std::cout << "thread exiting\n";
}

template <typename grid_type>
typename mcrx::grid_creator<grid_type>::refinement_accuracy_data
mcrx::grid_creator<grid_type>::recursive_refine_body (T_cell* c, 
						      int level,
						      std::vector<int>& stats )
{
  // Now check if we should refine the cell
  if (factory.refine_cell_p(*c, level)) {
    assert (c->is_leaf());// what would happen if this is not the case??
    c->refine();

    // and recursively refine those cells
    refinement_accuracy_data racc =
      recursive_refine (*c->sub_grid(), level + 1, stats);
    if (racc.all_leaves) {
      // recursive_refine returned true, which means that that grid
      // contains only leaf cells.  We can try to unrefine them
      unrefine_if_possible (*c, racc);
    }
    return racc;
  }
  else {
    // no, don't refine.  Get the cell data from the factory 
    if (stats.size() < level)
      stats.resize(level );
    ++stats [level - 1]; 
    c->set_data(new T_data (factory.get_data(*c)));
    return refinement_accuracy_data (*c->data (),*c->data()* *c->data(), 1, true);
  }
}


/** This function goes through the initial cells and subdivides them
    as they are put in the cell_stack if the estimated number of
    refinements needed is greater than work_chunk_levels.  This
    improves load balancing by ensuring that one thread doesn't get
    stuck with the core of the galaxy, taking many times longer than
    all the other threads. */
template <typename grid_type>
void
mcrx::grid_creator<grid_type>::create_balanced_queue (T_cell*c,
					      int level)
{
  if (factory.estimated_levels_needed () > factory.work_chunk_levels () ) {
    // This is a "high workload cell", subdivide it
    assert (c->is_leaf());
    c->refine();
    for (typename T_grid::local_iterator cc = c->sub_grid()->lbegin ();
	 cc != c->sub_grid()->lend (); ++cc)
      create_balanced_queue (cc, level + 1);
  }
  else {
    // a "low workload cell", put it in the cell_stack
    cell_stack.push_back(std::pair<T_cell*, int> (c, level) );
  }
}
				    
   
/** This function is responsible for actually building the grid.  Note
    that until we work out how to do multithreaded factory building,
    we can't run with more than one thread. It has something to do
    with the fact that inserting cells invalidates the list of cells
    the other threads should work on and makes things go bad. */
template <typename grid_type>
void mcrx::grid_creator<grid_type>::
build_grid (grid_type& g)
{
  int n_threads = factory.n_threads();
  if(n_threads>1) {
    std::cerr << "Error: Factory currently can't use >1 threads.\n";
    n_threads=1;
  }

  // put cells in cell stack
  if (work_chunk_levels > 0){
    std::cout << "creating balanced work chunks" << std::endl;
    for (typename T_grid::local_iterator i = g.lbegin ();
	 i != g.lend (); ++i)
      create_balanced_queue (i, 1);
  }
  else 
    for (typename T_grid::local_iterator i = g.lbegin (); 
	 i != g.lend (); ++i)
      cell_stack.push_back(std::pair<T_cell*, int> (i, 1) );
  
  std::cout << "beginning work on " << g.n_cells () << " cells" << std::endl;
  
  if (n_threads>1) {
    // spawn threads
    boost::thread_group threads;
    for (int i = 0; i < n_threads; ++i)
      //the thread_start object calls pop_cell_and_refine ()
      threads.create_thread(thread_start (this));
    threads.join_all();
  }
  else // no threads
    pop_cell_and_refine ();

  std::cout << "Final grid contains " << g.n_cells () << " cells" << std::endl;

}

#endif
